import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

# 加载数据
file_path = 'yiche_brands_sales.csv'
df = pd.read_csv(file_path)

# 处理数据，将日期转换为datetime格式
df['date'] = pd.to_datetime(df['date'])
df['timestamp'] = df['date'].map(pd.Timestamp.timestamp)

# 按品牌分组进行预测，这里以第一个品牌为例
brand_name = df['brand'].unique()[0]
brand_data = df[df['brand'] == brand_name]

# 准备特征和目标变量
X = brand_data[['timestamp']]
y = brand_data['sales']

# 使用多项式回归进行预测并绘制拟合曲线
poly = PolynomialFeatures(degree=3)
X_poly = poly.fit_transform(X)

# 划分训练集和测试集
X_train_poly, X_test_poly, y_train, y_test = train_test_split(X_poly, y, test_size=0.2, random_state=42)

# 多项式线性回归模型训练
model = LinearRegression()
model.fit(X_train_poly, y_train)

# 预测
y_pred = model.predict(X_test_poly)

# 计算均方误差
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

# 生成未来半年的日期
future_dates = pd.date_range(start=df['date'].max(), periods=180, freq='D')
future_timestamps = future_dates.map(pd.Timestamp.timestamp)
future_X_poly = poly.transform(future_timestamps.values.reshape(-1, 1))

# 预测未来销量
future_sales_pred = model.predict(future_X_poly)

# 可视化结果
plt.figure(figsize=(14, 7))
plt.scatter(brand_data['date'], brand_data['sales'], color='blue', label='Actual Sales')
plt.plot(future_dates, future_sales_pred, color='red', linewidth=2, label='Predicted Sales')
plt.xlabel('Date')
plt.ylabel('Sales')
plt.title(f'Sales Prediction for {brand_name}')
plt.legend()
plt.show()

print(f'Mean Squared Error: {mse}')
print(f'R^2 Score: {r2}')
